# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import numpy as np
import wrapt

from experimental.prophet.graph.shape import Shape
from experimental.prophet.graph.util import from_data_type_proto
from experimental.prophet.ops import constant
from experimental.prophet.ops import elemwise_math
from experimental.prophet.ops import cast


def as_variable(v):
  if isinstance(v, (Variable, NamedVariableRef)):
    return v
  else:
    return constant.constant(v)


class NamedVariableRef(wrapt.ObjectProxy):
  def __init__(self, name, variable):
    super().__init__(variable)
    assert name
    self._self_name = name

  def unref(self):
    return self.__wrapped__.unref()

  @property
  def name(self):
    return self._self_name

  @property
  def is_named_ref(self):
    return True


class Variable:
  def __init__(self, node, output_idx, arg_def):
    self._node = node
    self._output_idx = output_idx
    self._arg_def = arg_def
    self._dtype = from_data_type_proto(self.arg_def.dtype)
    if arg_def.attrs_json:
      self._attrs = json.loads(arg_def.attrs_json)
    else:
      self._attrs = {}
    self._arg_def_shape = Shape.from_proto(arg_def.shape)

    self._instance = None
    self._value_ref = None

  def unref(self):
    return self

  @property
  def is_constant(self):
    return self.node.is_constant

  @property
  def is_placeholder(self):
    return self.node.is_placeholder

  @property
  def name(self):
    return None

  @property
  def is_named_ref(self):
    return False

  def named_ref(self, name, namespace=None):
    if namespace is None:
      namespace = self.graph.building_ctx.get_namespace()
    name = namespace + '/' + name if namespace else name
    var_ref = NamedVariableRef(name, self)
    self.graph.building_ctx.register_named_variable_ref(var_ref)
    return var_ref

  @property
  def node(self):
    return self._node

  @property
  def output_idx(self):
    return self._output_idx

  @property
  def arg_def(self):
    return self._arg_def

  @property
  def dtype(self):
    return self._dtype

  @property
  def graph(self):
    return self.node.graph

  @property
  def arg_name(self):
    return self.arg_def.name

  @property
  def attrs(self):
    return self._attrs

  @property
  def shape(self):
    if not self.is_instantiated or self.is_null:
      return self._arg_def_shape
    else:
      return self._instance.shape

  @property
  def is_scalar(self):
    return Shape(self.shape).is_scalar

  @property
  def value(self):
    if self._instance is None:
      self.instantiate()
    return self._get_value()

  def _get_value(self):
    if self._instance.is_null():
      return None

    if self._value_ref is None or self._instance.is_buffer_updated():
      self._value_ref = np.array(self._instance, copy=False)

    if len(self._instance.shape) == 0:
      return self._value_ref.item()
    else:
      return self._value_ref

  @value.setter
  def value(self, value):
    if value is None:
      self.instance.set_null(True)
    else:
      self.instance.set_value(value)  # will set_null(False) internally.

  @property
  def is_null(self):
    return self.instance.is_null()

  @property
  def is_instantiated(self):
    try:
      self.instance
      return True
    except Exception:
      return False

  def instantiate(self):
    if self._instance is None:
      self._instance = self.graph._get_variable_instance(self.instance_name)

  @property
  def instance_name(self):
    return '%s:%s' % (self.node.name, self.arg_def.name)

  @property
  def instance(self):
    if self._instance is None:
      self.instantiate()
    return self._instance

  def _reset_instance(self):
    self._instance = None
    self._value_ref = None

  def cast(self, dtype):
    return cast.cast(self, dtype)

  # Operator overloading
  # Comparison operators
  def __lt__(self, other):
    return elemwise_math.less(self, as_variable(other))

  def __le__(self, other):
    return elemwise_math.less_equal(self, as_variable(other))

  def __eq__(self, other):
    return elemwise_math.equal(self, as_variable(other))

  def __ne__(self, other):
    return elemwise_math.not_equal(self, as_variable(other))

  def __gt__(self, other):
    return elemwise_math.greater(self, as_variable(other))

  def __ge__(self, other):
    return elemwise_math.greater_equal(self, as_variable(other))

  # Bitwise operators
  def __invert__(self):
    assert self.dtype in (bool, int)
    return elemwise_math.bitwise_not(self)

  def __lshift__(self, other):
    assert self.dtype in (bool, int)
    return elemwise_math.bitwise_left_shift(self, as_variable(other))

  def __rlshift__(self, other):
    assert self.dtype in (bool, int)
    return elemwise_math.bitwise_left_shift(as_variable(other), self)

  def __rshift__(self, other):
    assert self.dtype in (bool, int)
    return elemwise_math.bitwise_right_shift(self, as_variable(other))

  def __rrshift__(self, other):
    assert self.dtype in (bool, int)
    return elemwise_math.bitwise_right_shift(as_variable(other), self)

  def __xor__(self, other):
    return elemwise_math.bitwise_xor(self, as_variable(other))

  def __rxor__(self, other):
    return elemwise_math.bitwise_xor(as_variable(other), self)

  def __and__(self, other):
    return elemwise_math.bitwise_and(self, as_variable(other))

  def __rand__(self, other):
    return elemwise_math.bitwise_and(as_variable(other), self)

  def __or__(self, other):
    return elemwise_math.bitwise_or(self, as_variable(other))

  def __ror__(self, other):
    return elemwise_math.bitwise_or(as_variable(other), self)

  # Logical operators
  # Logical operator overloading is not supported in Pythhon.
  # Refer to https://www.python.org/dev/peps/pep-0335

  # Numeric operators
  def __neg__(self):
    return elemwise_math.neg(self)

  def __pos__(self):
    return self

  def __abs__(self):
    return elemwise_math.abs(self)

  def __add__(self, other):
    return elemwise_math.add(self, as_variable(other))

  def __radd__(self, other):
    return elemwise_math.add(as_variable(other), self)

  def __sub__(self, other):
    return elemwise_math.sub(self, as_variable(other))

  def __rsub__(self, other):
    return elemwise_math.sub(as_variable(other), self)

  def __mul__(self, other):
    return elemwise_math.mul(self, as_variable(other))

  def __rmul__(self, other):
    return elemwise_math.mul(as_variable(other), self)

  def __truediv__(self, other):
    return elemwise_math.div(self, as_variable(other))

  def __rtruediv__(self, other):
    return elemwise_math.div(as_variable(other), self)

  def __mod__(self, other):
    other = as_variable(other)
    if self.dtype in (bool, int) and other.dtype in (bool, int):
      return elemwise_math.mod(self, as_variable(other))
    else:
      return elemwise_math.fmod(self, as_variable(other))

  def __pow__(self, other):
    return elemwise_math.pow(self, as_variable(other))

  def __rpow__(self, other):
    return elemwise_math.pow(as_variable(other), self)

  def __round__(self):
    return elemwise_math.round(self, as_variable(other))

  def __trunc__(self):
    return elemwise_math.trunc(self, as_variable(other))

  def __floor__(self):
    return elemwise_math.floor(self, as_variable(other))

  def __ceil__(self):
    return elemwise_math.ceil(self, as_variable(other))

  # TODO(jaewon): Imeplement __i*__ methods, __iadd__, __isub__, ...
